#include <stdio.h>
#include "packet_capture.h"


// static int get_nic_hwaddr(int fd, const char *dev_name)
// {
//     struct ifreq ifr;
//     bzero(&ifr, sizeof(ifr));
//     strncpy(ifr.ifr_name, dev_name, IFNAMSIZ);
// 
//     /* SIOCGIFHWADDR: get NIC`s hardware address */
//     if (ioctl(fd, SIOCGIFHWADDR, &ifr) == -1)
//     {
//         printf("Error---->ioctl(SIOCGIFHWADDR) error: %s\n", strerror(errno));
//         return AF_FAILED;
//     }
// #if 0
//     if (strcmp(dev_name, g_dev_info.name) == 0)
//     {
//         memcpy(g_dev_info.mac, ifr.ifr_hwaddr.sa_data, ETH_ALEN);
//     }
//     else
//     {
//         // todo
//     }
// #endif
// 
//     return AF_SUCCEED;
// }

static int get_nic_index(int fd, const char *dev_name)
{
    struct ifreq ifr;
    bzero(&ifr, sizeof(ifr));
    strncpy(ifr.ifr_name, dev_name, IFNAMSIZ);

    /* SIOCGIFINDEX: get NIC`s index  */
    if (ioctl(fd, SIOCGIFINDEX, &ifr) == -1)
    {
        printf("Error---->ioctl(SIOCGIFINDEX) error: %s\n", strerror(errno));
        return AF_FAILED;
    }

    return ifr.ifr_ifindex;
}


static void destroy_instance(afpacket_instance *p_ins)
{
    if (p_ins == NULL)
        return;

    u_int16_t ring_size = 0;
    struct tpacket_req req;
    if (p_ins->fd != -1)
    {
        /* destroy userspace RX ring */
        if (p_ins->rx_ring.entry != NULL)
        {
            free(p_ins->rx_ring.entry);
            p_ins->rx_ring.entry = NULL;
        }
        /* destroy userspace TX ring */
        if (p_ins->tx_ring.entry != NULL)
        {
            free(p_ins->tx_ring.entry);
            p_ins->tx_ring.entry = NULL;
        }
        /* unmmap kernel packet ring */
        if (p_ins->buffer != MAP_FAILED)
        {
            ring_size = p_ins->rx_ring.size + p_ins->tx_ring.size;
            munmap(p_ins->buffer, ring_size);
            p_ins->buffer = MAP_FAILED;
        }

        /* tell kernel to destroy rings */
        bzero(&req, sizeof(req));
        setsockopt(p_ins->fd, SOL_PACKET, PACKET_RX_RING, (void *)&req, sizeof(req));
        setsockopt(p_ins->fd, SOL_PACKET, PACKET_TX_RING, (void *)&req, sizeof(req));
        close(p_ins->fd);
    }

    if (p_ins->name != NULL)
    {
        free(p_ins->name);
        p_ins->name = NULL;
    }

    free(p_ins);
    p_ins = NULL;
}


static afpacket_instance *create_instance(const char *dev_name)
{
    if (dev_name == NULL)
    {
        printf("Error----> dev_name is null !\n");
        return NULL;
    }
    afpacket_instance *p_ins = (afpacket_instance *)calloc(1, sizeof(afpacket_instance));
    if (p_ins == NULL)
    {
        printf("Error---->[%s] could not allocate a new instance structure .\n", __FUNCTION__);
        goto err;
    }
    if ((p_ins->name = strdup(dev_name)) == NULL)
    {
        printf("Error---->[%s] could not allocate a copy instance structure .\n", __FUNCTION__);
        goto err;
    }

    // create raw socket
    p_ins->fd = socket(PF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
    if (p_ins->fd == -1)
    {
        printf("Error---->[%s] could not open AF_PACKET socket .\n", __FUNCTION__);
        goto err;
    }

    // // get nic`s hw_addr
    // if (get_nic_hwaddr(p_ins->fd, p_ins->name) != AF_SUCCEED)
    // {
    //     printf("Error---->[%s] could not get hardware address for nic .\n", __FUNCTION__);
    //     goto err;
    // }

    // get nis`s index
    p_ins->index = get_nic_index(p_ins->fd, p_ins->name);
    if (p_ins->index == -1)
    {
        printf("Error---->[%s] could not find index for nic .\n", __FUNCTION__);
        goto err;
    }
    printf("------->p_ins->index = %d, p_ins->name = %s\n", p_ins->index, p_ins->name);

    // set sll attribution
    p_ins->sll.sll_family = AF_PACKET;
    p_ins->sll.sll_ifindex = p_ins->index;
    p_ins->sll.sll_protocol = htons(ETH_P_ALL);
    p_ins->ring_size = RING_SIZE * MAX_NUM;
    p_ins->snaplen = AF_MTU;

    return p_ins;

err:
    destroy_instance(p_ins);
    return NULL;
}

static int bind_interface(afpacket_instance *p_ins)
{
    if (bind(p_ins->fd, (struct sockaddr *)&(p_ins->sll), sizeof(p_ins->sll)) != 0)
    {
        printf("Error==>[%s] bind interface failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* check pending errors */
    int err = 0;
    socklen_t errlen = sizeof(err);
    if (getsockopt(p_ins->fd, SOL_SOCKET, SO_ERROR, &err, &errlen) || err)
    {
        printf("Error==>[%s] getsockopt(SO_ERROR) failed: %s\n", __FUNCTION__, strerror(errno));
        return AF_FAILED;
    }

    return AF_SUCCEED;
}

static int set_nic_promisc(afpacket_instance *p_ins)
{
    struct packet_mreq mr;
    bzero(&mr, sizeof(mr));
    mr.mr_ifindex = p_ins->index;
    mr.mr_type = PACKET_MR_PROMISC;
    if (setsockopt(p_ins->fd, SOL_PACKET, PACKET_ADD_MEMBERSHIP, &mr, sizeof(mr)) == -1)
    {
        printf("Error==>[%s] setsockopt(PACKET_ADD_MEMBERSHIP) failed: %s\n", __FUNCTION__, strerror(errno));
        return AF_FAILED;
    }

    return AF_SUCCEED;
}

static int get_arptype(afpacket_instance *p_ins)
{
    struct ifreq ifr;
    bzero(&ifr, sizeof(ifr));
    strncpy(ifr.ifr_name, p_ins->name, sizeof(ifr.ifr_name));
    if (ioctl(p_ins->fd, SIOCGIFHWADDR, &ifr) == -1)
    {
        printf("Error==>[%s] ioctl(SIOCGIFHWADDR) failed: %s\n", __FUNCTION__, strerror(errno));
        return AF_FAILED;
    }

    return ifr.ifr_hwaddr.sa_family;
}

static int determine_version(afpacket_instance *p_ins)
{
    /* Probe whether kernel supports TPACCKET_V2 */
    int val = TPACKET_V2;
    socklen_t len = sizeof(val);
    if (getsockopt(p_ins->fd, SOL_PACKET, PACKET_HDRLEN, &val, &len) < 0)
    {
        return AF_FAILED;
    }
    p_ins->tp_hdrlen = val;

    /* tell kernel to use TPACKET_V2 */
    val = TPACKET_V2;
    if (setsockopt(p_ins->fd, SOL_PACKET, PACKET_VERSION, &val, sizeof(val)) < 0)
    {
        printf("Error==>[%s] setsockopt(PACKET_VERSION) failed: %s\n", __FUNCTION__, strerror(errno));
        return AF_FAILED;
    }
    p_ins->tp_version = TPACKET_V2;

    /* reserve space for vlan tag reconstruction */
    val = VLAN_TAG_LEN;
    if (setsockopt(p_ins->fd, SOL_PACKET, PACKET_RESERVE, &val, sizeof(val)) < 0)
    {
        printf("Error==>[%s] setsockopt(PACKET_RESERVE) failed: %s\n", __FUNCTION__, strerror(errno));
        return AF_FAILED;
    }

    return AF_SUCCEED;
}

static int calculate_layout(afpacket_instance *p_ins, struct tpacket_req *layout, u_int32_t tp_hdrlen, u_int32_t order)
{
    u_int32_t tp_hdrlen_sll = TPACKET_ALIGN(tp_hdrlen) + sizeof(struct sockaddr_ll);
    u_int32_t netoff = TPACKET_ALIGN(tp_hdrlen_sll + ETH_HLEN) + VLAN_TAG_LEN;
    layout->tp_frame_size = TPACKET_ALIGN(netoff - ETH_HLEN + p_ins->snaplen);
    layout->tp_block_size = getpagesize() << order;

    while (layout->tp_block_size < layout->tp_frame_size)
    {
        layout->tp_block_size <<= 1;
    }

    u_int32_t frames_per_block = layout->tp_block_size / layout->tp_frame_size;
    if (frames_per_block == 0)
    {
        printf("Error==>[%s] invalid frames per block !\n", __FUNCTION__);
        return AF_FAILED;
    }

    layout->tp_frame_nr = p_ins->ring_size / layout->tp_frame_size;
    layout->tp_block_nr = layout->tp_frame_nr / frames_per_block;
    layout->tp_frame_nr = layout->tp_block_nr * frames_per_block;

    return AF_SUCCEED;
}

static int create_ring(afpacket_instance *p_ins, afpacket_ring *ring, int optname)
{
    u_int32_t order = 0;
    /* start with page allocation of order 3, try to allocate RX ring in kernel */
    for (order=DEFAULT_ORDER; order >= 0; order--)
    {
        if (calculate_layout(p_ins, &ring->layout, p_ins->tp_hdrlen, order) != AF_SUCCEED)
        {
            return AF_SUCCEED;
        }

        /* ask kernel to create ring */
        if (setsockopt(p_ins->fd, SOL_PACKET, optname, (void *)&ring->layout, sizeof(struct tpacket_req)) != AF_SUCCEED)
        {
            if (errno == ENOMEM)
            {
                printf("Error==>%s: allocation of kernel packet ring failed with order %d, retry ...\n", p_ins->name, order);
                continue;
            }
            return AF_FAILED;
        }

        /* reserve total ring size for later */
        ring->size = (ring->layout.tp_block_size * ring->layout.tp_block_nr);
        printf("tp_block_size=%d, tp_block_nr=%d, ring->size=%d\n", ring->layout.tp_block_size, ring->layout.tp_block_nr, ring->size);

        return AF_SUCCEED;
    }
    return AF_FAILED;
}

static int mmap_ring(afpacket_instance *p_ins)
{
    /* mmap the ring into userspace */
    u_int32_t ring_size = p_ins->rx_ring.size + p_ins->tx_ring.size;
    // printf("rx.size=%d, tx.size=%d\n", p_ins->rx_ring.size, p_ins->tx_ring.size);
    p_ins->buffer = mmap(NULL, ring_size, PROT_READ | PROT_WRITE, MAP_SHARED, p_ins->fd, 0);
    if (p_ins->buffer == MAP_FAILED)
    {
        printf("Error==>[%s] could not MMAP ring: %s\n", __FUNCTION__, strerror(errno));
        printf("------->buffer=%p, ring_size=%d, fd=%d\n", p_ins->buffer, ring_size, p_ins->fd);
        return AF_FAILED;
    }

    p_ins->rx_ring.start = (u_int8_t *)p_ins->buffer;
    p_ins->tx_ring.start = (u_int8_t *)p_ins->buffer + p_ins->rx_ring.size; 

    return AF_SUCCEED;
}

static int setup_ring(afpacket_ring *ring)
{
    u_int32_t idx, block, block_offset, frame, frame_offset;

    /* allocate a ring to hold packet pointer */
    ring->entry = (afpacket_entry *)calloc(ring->layout.tp_frame_nr, sizeof(afpacket_entry));
    if (ring->entry == NULL)
    {
        printf("Error==>[%s] could not allocate ring\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* setup the buffer entry pointers in the ring */
    for (idx=0, block=0; block < ring->layout.tp_block_nr; block++)
    {
        block_offset = block * ring->layout.tp_block_size;
        for (frame=0; 
                frame < (ring->layout.tp_block_size/ring->layout.tp_frame_size)
                && idx < ring->layout.tp_frame_nr; frame++)
        {
            frame_offset = frame * ring->layout.tp_frame_size;
            ring->entry[idx].hdr.raw = (u_int8_t *)ring->start + block_offset + frame_offset;
            ring->entry[idx].next = &(ring->entry[idx + 1]);
            idx++;
        }
    }

    /* make this become a circular buffer */
    ring->entry[ring->layout.tp_frame_nr - 1].next = &ring->entry[0];

    /* initialize entry point into ring as first buffer entry */
    ring->cursor = &ring->entry[0];

    return AF_SUCCEED;
}

static void reset_statistics(afpacket_instance *p_ins)
{
    struct tpacket_stats kstats;
    bzero(&kstats, sizeof(kstats));
    bzero(&p_ins->stats, sizeof(afpacket_statistics));
    socklen_t len = sizeof(struct tpacket_stats);
    getsockopt(p_ins->fd, SOL_PACKET, PACKET_STATISTICS, &kstats, &len);
}


/* -------------------------------------------------------------------------- */
int afpacket_init(const char *dev_name, void **context_ptr)
{
    afpacket_instance *p_ins = NULL;
    p_ins = create_instance(dev_name);
    if (p_ins == NULL)
    {
        return AF_FAILED;
    }
    *context_ptr = p_ins;
    return AF_SUCCEED;
}

int afpacket_start(void *ins, u_int8_t flag)
{
    /* flag: 0:send, 1:recv */
    afpacket_instance *p_ins = (afpacket_instance *)ins;
    if (bind_interface(p_ins) != AF_SUCCEED)
    {
        printf("Error==>[%s] bind failed !\n", __FUNCTION__);
        return AF_FAILED;
    }
    if (flag == 0)
        return AF_SUCCEED;

    if (set_nic_promisc(p_ins) != AF_SUCCEED)
    {
        printf("Error==>[%s] set nic promisc failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* get link-level type */
    int arp_type = get_arptype(p_ins);
    if (arp_type < 0)
    {
        printf("Error==>[%s] get arp type failed !\n", __FUNCTION__);
        return AF_FAILED;
    }
    if (arp_type != ARPHRD_ETHER)
    {
        printf("Error==>[%s] arp_type != ARPHRD_ETHER\n", __FUNCTION__);
        return AF_FAILED;
    }

    if (determine_version(p_ins) != AF_SUCCEED)
    {
        printf("Error==>[%s] determine version failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* create ring */
    if (create_ring(p_ins, &p_ins->rx_ring, PACKET_RX_RING) != AF_SUCCEED)
    {
        printf("Error==>[%s] create rx_ring failed !\n", __FUNCTION__);
        return AF_FAILED;
    }
    if (create_ring(p_ins, &p_ins->tx_ring, PACKET_TX_RING) != AF_SUCCEED)
    {
        printf("Error==>[%s] create tx_ring failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* mmap ring */
    if (mmap_ring(p_ins) != AF_SUCCEED)
    {
        printf("Error==>[%s] mmap ring failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* setup ring */
    if (setup_ring(&p_ins->rx_ring) != AF_SUCCEED)
    {
        printf("Error==>[%s] setup rx_ring failed !\n", __FUNCTION__);
        return AF_FAILED;
    }
    if (setup_ring(&p_ins->tx_ring) != AF_SUCCEED)
    {
        printf("Error==>[%s] setup tx_ring failed !\n", __FUNCTION__);
        return AF_FAILED;
    }

    /* reset packet statistics */
    reset_statistics(p_ins);

    return AF_SUCCEED;
}

int afpacket_poll(void *ins, yx_pkt *pkt, u_int32_t len)
{
    afpacket_instance *p_ins = (afpacket_instance *)ins;
    u_int32_t tp_mac = 0;
    u_int32_t tp_snaplen = 0;
    const u_int8_t *p_data;
    union thdr tp_hdr = p_ins->rx_ring.cursor->hdr;
    if (tp_hdr.h2->tp_status & TP_STATUS_USER)
    {
        tp_mac = tp_hdr.h2->tp_mac;
        tp_snaplen = tp_hdr.h2->tp_snaplen;

        /* tpacket mac tp_snaplen check */
        if (tp_snaplen == 0)
        {
            goto RING_NEXT;
        }
        if (tp_mac + tp_snaplen > p_ins->rx_ring.layout.tp_frame_size)
        {
            tp_snaplen = 0;
            goto RING_NEXT;
        }
        if (tp_hdr.h2->tp_vlan_tci && AF_VLAN_ENABLE)
        {
        }
        if (tp_snaplen > len)
        {
            tp_snaplen = len;
        }

        p_data = p_ins->rx_ring.cursor->hdr.raw + tp_mac;
        bzero(pkt->rx->buf, AF_MTU);
        memcpy(pkt->rx->buf, p_data, tp_snaplen);
        pkt->rx->len = tp_snaplen;

RING_NEXT:
        tp_hdr.h2->tp_status = TP_STATUS_KERNEL;
        p_ins->rx_ring.cursor = p_ins->rx_ring.cursor->next;

        return tp_snaplen;
    }

    struct pollfd pfd;
    bzero(&pfd, sizeof(pfd));
    pfd.fd = p_ins->fd;
    pfd.events = POLLIN;
    pfd.revents = 0;
    while (poll(&pfd, 1, INFTIM) <= 0);

    return 0;
}

int afpacket_send(void *ins, yx_pkt *pkt)
{
    afpacket_instance *p_ins = (afpacket_instance *)ins;

    // 14 + 20 + 20/8 + 4 + 6/18 = 64byte
    // ether + ip + tcp/udp + fcs + padding = 64byte
    if (pkt->tx->len < 60)
    {
        pkt->tx->len = 60;
    }

    int ret = send(p_ins->fd, pkt->tx->buf, pkt->tx->len, 0);

    return ret;
}

int afpacket_close(void *ins)
{
    afpacket_instance *p_ins = (afpacket_instance *)ins;
    destroy_instance(p_ins);
    return AF_SUCCEED;
}
